import numpy as np
from torch_geometric.data import Data
import enum
import torch
from torch_scatter import scatter


class MeshType(enum.IntEnum):
    Triangle = 1
    Tetrahedron = 2
    Quad = 3
    Line = 4
    Flat = 5


def calc_cell_centered_with_node_attr(
    node_attr, cells_node, cells_index, reduce="mean", map=True
):
    if cells_node.shape != cells_index.shape:
        raise ValueError("wrong cells_node/cells_index dim")

    if len(cells_node.shape) > 1:
        cells_node = cells_node.view(-1)

    if len(cells_index.shape) > 1:
        cells_index = cells_index.view(-1)

    if map:
        mapped_node_attr = node_attr[cells_node]
    else:
        mapped_node_attr = node_attr

    cell_attr = scatter(src=mapped_node_attr, index=cells_index, dim=0, reduce=reduce)

    return cell_attr


def calc_node_centered_with_cell_attr(
    cell_attr, cells_node, cells_index, reduce="mean", map=True
):
    if cells_node.shape != cells_index.shape:
        raise ValueError(f"wrong cells_node/cells_index dim ")

    if len(cells_node.shape) > 1:
        cells_node = cells_node.view(-1)

    if len(cells_index.shape) > 1:
        cells_index = cells_index.view(-1)

    if map:
        maped_cell_attr = cell_attr[cells_index]
    else:
        maped_cell_attr = cell_attr

    cell_attr = scatter(src=maped_cell_attr, index=cells_node, dim=0, reduce=reduce)

    return cell_attr


# see https://github.com/sungyongs/dpgn/blob/master/utils.py
def decompose_and_trans_node_attr_to_cell_attr_graph(
    graph, has_changed_node_attr_to_cell_attr
):
    # graph: torch_geometric.data.data.Data
    # TODO: make it more robust
    x, edge_index, edge_attr, face, global_attr, ball_edge_index = (
        None,
        None,
        None,
        None,
        None,
        None,
    )

    for key in graph.keys():
        if key == "x":
            x = graph.x  # avoid exception
        elif key == "edge_index":
            edge_index = graph.edge_index
        elif key == "edge_attr":
            edge_attr = graph.edge_attr
        elif key == "global_attr":
            global_attr = graph.global_attr
        elif key == "face":
            face = graph.face
        elif key == "ball_edge_index":
            ball_edge_index = graph.ball_edge_index
        else:
            pass

    return (x, edge_index, edge_attr, face, global_attr, ball_edge_index)


# see https://github.com/sungyongs/dpgn/blob/master/utils.py
def copy_geometric_data(graph, has_changed_node_attr_to_cell_attr):
    """return a copy of torch_geometric.data.data.Data
    This function should be carefully used based on
    which keys in a given graph.
    """
    (
        node_attr,
        edge_index,
        edge_attr,
        face,
        global_attr,
        ball_edge_index,
    ) = decompose_and_trans_node_attr_to_cell_attr_graph(
        graph, has_changed_node_attr_to_cell_attr
    )

    ret = Data(
        x=node_attr,
        edge_index=edge_index,
        edge_attr=edge_attr,
        face=face,
        ball_edge_index=ball_edge_index,
    )

    ret.global_attr = global_attr

    return ret


def shuffle_np(array):
    array_t = array.copy()
    np.random.shuffle(array_t)
    return array_t
